import xgboost as xgb
import numpy as np
from sklearn.metrics import ndcg_score

"""
XGBoost 排序模型示例
排序问题常用于搜索引擎、推荐系统等场景
目标：学习如何对项目进行排序，而不是分类或回归
"""


def create_ranking_dataset():
    """
    创建排序数据集示例
    模拟搜索引擎场景：不同query下的文档排序
    """
    # 特征：文档质量、相关性、点击率等
    X = np.array(
        [
            # Query 1 的文档 (3个文档)
            [0.8, 0.9, 0.7],  # 高质量相关文档
            [0.6, 0.7, 0.5],  # 中等质量文档
            [0.3, 0.4, 0.2],  # 低质量文档
            # Query 2 的文档 (4个文档)
            [0.9, 0.8, 0.6],  # 高质量相关文档
            [0.7, 0.6, 0.8],  # 中等质量文档
            [0.5, 0.5, 0.4],  # 低质量文档
            [0.2, 0.3, 0.1],  # 最低质量文档
            # Query 3 的文档 (3个文档)
            [0.85, 0.75, 0.8],  # 高质量相关文档
            [0.6, 0.5, 0.7],  # 中等质量文档
            [0.4, 0.3, 0.2],  # 低质量文档
        ]
    )

    # 相关性标签：0-不相关，1-相关，2-高度相关
    y = np.array([2, 1, 0, 2, 1, 0, 0, 2, 1, 0])

    # 每个query的文档数量
    groups = np.array([3, 4, 3])

    return X, y, groups


def train_ranking_model(X, y, groups, max_depth=6, n_estimators=100, learning_rate=0.1):
    """
    训练XGBoost排序模型
    :param X: 特征矩阵
    :param y: 相关性标签
    :param groups: 每个query的文档数量
    :param max_depth: 树的最大深度
    :param n_estimators: 树的数量
    :param learning_rate: 学习率
    :return: 训练好的模型
    """
    # 创建DMatrix并设置组信息
    dtrain = xgb.DMatrix(X, label=y)
    dtrain.set_group(groups)

    # 排序模型参数
    params = {
        "objective": "rank:pairwise",  # 成对排序目标
        "max_depth": max_depth,  # 树的最大深度
        "eta": learning_rate,  # 学习率
        "gamma": 1.0,  # 正则化项系数
        "min_child_weight": 0.1,  # 叶子节点最小权重
        "subsample": 0.8,  # 子样本比例
        "eval_metric": "ndcg@3",  # 使用NDCG@3作为评估指标
    }

    # 训练模型
    model = xgb.train(
        params,
        dtrain,
        num_boost_round=n_estimators,
        evals=[(dtrain, "train")],
        verbose_eval=False,
    )

    return model


def evaluate_ranking_model(model, X, y, groups):
    """
    简化版排序模型评估
    :param model: 训练好的模型
    :param X: 特征矩阵
    :param y: 真实相关性标签
    :param groups: 每个query的文档数量
    :return: 平均NDCG@3
    """
    # 预测排序得分
    dtest = xgb.DMatrix(X)
    dtest.set_group(groups)
    predictions = model.predict(dtest)

    # 计算平均NDCG@3
    ndcg_scores = []
    start_idx = 0

    for group_size in groups:
        end_idx = start_idx + group_size
        group_y_true = y[start_idx:end_idx].reshape(1, -1)
        group_y_pred = predictions[start_idx:end_idx].reshape(1, -1)

        k = min(3, group_size)
        ndcg = ndcg_score(group_y_true, group_y_pred, k=k)
        ndcg_scores.append(ndcg)
        start_idx = end_idx

    mean_ndcg = np.mean(ndcg_scores)

    # 简洁输出
    print(f"排序模型评估 - 平均NDCG@3: {mean_ndcg:.4f}")
    return mean_ndcg, predictions


def save_model(model, model_path):
    """
    保存模型到指定路径
    :param model: 训练好的模型实例
    :param model_path: 模型保存路径
    """
    model.save_model(model_path)
    print(f"模型已保存至 {model_path}")


def load_model(model_path):
    """
    从指定路径加载模型
    :param model_path: 模型保存路径
    :return: 加载的模型实例
    """
    # 创建一个空的Booster对象
    model = xgb.Booster()
    model.load_model(model_path)
    print(f"模型已从 {model_path} 加载")
    return model


def predict(model, X):
    """
    使用模型进行预测
    :param model: 训练好的模型实例
    :param X: 输入特征数据
    :return: 预测结果
    """
    return model.predict(X)


# 主程序执行
if __name__ == "__main__":
    # 创建数据集
    print("创建排序数据集...")
    X, y, groups = create_ranking_dataset()

    print(f"数据集信息:")
    print(f"  样本数: {len(X)}")
    print(f"  特征数: {X.shape[1]}")
    print(f"  Query数量: {len(groups)}")
    print(f"  各组文档数: {groups}")
    print(f"  相关性标签分布: {np.bincount(y)}")

    # 训练模型
    print("\n开始训练排序模型...")
    model = train_ranking_model(X, y, groups)

    # 评估模型
    print("\n评估模型性能...")
    mean_ndcg, predictions = evaluate_ranking_model(model, X, y, groups)

    # 保存模型
    # save_model(model, "xgboost_ranking_model.json")

    # 加载模型示例
    # print("\n验证模型加载功能:")
    # loaded_model = load_model("xgboost_ranking_model.json")

    print("\n程序执行完成！")
